package beast.evolution.tree;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Random;

import beast.core.Description;
import beast.core.Input;
import beast.evolution.alignment.Alignment;
import beast.evolution.datatype.Nucleotide;

@Description("This class extends Tree and create all nodes as Node2, so that sequence is part of the state")
public class Tree3 extends Tree {
	public Input<Alignment> m_alignment = new Input<Alignment>("alignment",
			"alignment that corresponds to the sequence data");

	int[][] sequences;

	int sequenceSize;

	// List to store pMatrix
	List<double[][]> pMatrixList = new ArrayList<double[][]>();

	// List to store Char state
	List<int[]> charStateList = new ArrayList<int[]>();

	// random number generator
	Random generator = new Random();

	@Override
	public void initAndValidate() throws Exception {
		if ((m_initial != null) && (m_alignment != null)) {
			super.initAndValidate();
			sequences = initSequences();
			sequenceSize = m_alignment.get().getSiteCount();

			// add same number of pMatrix as node number to pMatrixList
			for (int i = 0; i < m_nodes.length; i++) {
				pMatrixList.add(new double[4][4]);
				charStateList.add(new int[] { -1, -1, -1, -1 });
			}

			// temporarily commented out for debugging purpose
			// PupkoAllSites();

			showSequences();
		}
	}

	public void showSequences() {
		Nucleotide nucleo = new Nucleotide();
		for (int i = 0; i < m_nodes.length; i++) {
			int[] tmp_seq = sequences[i];
			/**
			 * System.out.println("The length of the sequence is:" +
			 * aligned_seq.getSiteCount());
			 */
			String real_seq = nucleo.state2string(tmp_seq);
			System.out.println("The node is:" + i);
			if (m_nodes[i].isLeaf()) {
				System.out.println("The ID is:" + m_nodes[i].getID()
						+ " and length is: " + m_nodes[i].getLength());
			}
			System.out.println(" and the Seq is: \n" + real_seq);
		}
	}

	/**
	 * Assign initial values of sequence info to Leaf nodes based on Alignment
	 * Assign zero values (all As) to internal nodes and root
	 */
	public int[][] initSequences() throws Exception {
		int[][] allSequences = new int[nodeCount][m_alignment.get()
				.getSiteCount()];

		/* for leaf sequences, we assign values based on taxon match */
		for (int i = 0; i < getLeafNodeCount(); i++) {
			String sequenceID = getNode(i).getID();
			allSequences[i] = getSequenceByID(sequenceID);
			System.out.println(Arrays.toString(allSequences[i]));
		}

		return allSequences;
	}

	public int[] getSequenceByID(String sequenceID) throws Exception {
		Nucleotide nucleo = new Nucleotide();
		List<Integer> sequenceTarget = new ArrayList<Integer>();
		for (int i = 0; i < m_alignment.get().getNrTaxa(); i++) {
			if (m_alignment.get().m_pSequences.get().get(i).m_sTaxon.get()
					.toString() == sequenceID) {
				sequenceTarget = m_alignment.get().m_pSequences.get().get(i)
						.getSequence(nucleo);
			}
		}
		int[] sequence = new int[sequenceTarget.size()];
		for (int i = 0; i < sequence.length; i++) {
			sequence[i] = sequenceTarget.get(i);
		}
		System.out.println(Arrays.toString(sequence));
		return sequence;
	}

	public int[] getSequenceByNr(int nodeNum) {
		int[] sequence = sequences[nodeNum];
		return sequence;
	}

	// add Pupko to Tree3 (substitute corresponding parts in Tree2)
	public void PupkoAllSites() {
		// apply PupkoOneSite to all sites
		for (int i = 0; i < sequenceSize; i++) {
			PupkoOneSite(i);
		}
	}

	// the Pupko algorithm
	public void PupkoOneSite(int siteNum) {
		/*
		 * temp: an array to temp store new char state Left to right: A C G T =>
		 * 0 1 2 3
		 */
		int[] newCharState = new int[] { 1, 1, 1, 1 };

		// create an array to store nucleotide states at the temp site for each
		// node
		int[] m_states = new int[m_nodes.length];
		for (int i = 0; i < m_nodes.length; i++) {
			m_states[i] = sequences[i][siteNum];
			// System.out.println("The node is:" + m_nodes[i].getNr()
			// + " and the nucleotide state is" + m_states[i]);
		}

		/* deal with leaf nodes first (since they can be readily calculated) */
		for (int i = 0; i < root.getLeafNodeCount(); i++) {
			/*
			 * given the observed nucleo type j at leaf node, set the jth column
			 * of the pMatrix
			 */
			setLeafpMatrix(m_states[i], pMatrixList.get(i), m_nodes[i]
					.getLength());
			/*
			 * given the observed nucleo type at leaf node i, we determine
			 * "newCharState" for this leaf node and update the CharState of
			 * this leaf node
			 */
			setLeafCharState(i, newCharState, m_states);

			System.out.println("The species name is:"
					+ m_nodes[i].getID()
					+ "\n The node is:"
					+ m_nodes[i].getNr()
					+ " and the nucleotide state is"
					+ Arrays.toString(newCharState)
					+ "\n and the branch length is:"
					+ m_nodes[i].getLength()
					+ " and the pMatrix is: \n"
					+ Arrays.deepToString(pMatrixList.get(i)).replace("], ",
							"\n").replaceAll(",|\\[|\\]", ""));
		}

		// create a list to store numbers for internal nodes
		List<Integer> internalNodeNum = new ArrayList<Integer>();
		for (int i = root.getLeafNodeCount(); i < root.getNodeCount(); i++) {
			internalNodeNum.add(i);
		}
		// remove the root number from internalNodeNum
		internalNodeNum.remove(internalNodeNum.indexOf(root.getNr()));

		/*
		 * while internalNodeNum is not empty, constantly search through
		 * internal node list and find those nodes whose children have already
		 * been calculated
		 */

		while (internalNodeNum.size() > 0) {
			for (int i = 0; i < internalNodeNum.size(); i++) {
				if (m_nodes[internalNodeNum.get(i)].getLength() > 0) {
					// get the node number of the left child node and right
					// child
					// node
					int leftNum = m_nodes[internalNodeNum.get(i)].m_left
							.getNr();
					int rightNum = m_nodes[internalNodeNum.get(i)].m_right
							.getNr();
					// if both left and right child node has been updated, then
					// update this internal node
					if ((charStateList.get(leftNum)[0] != -1)
							&& (charStateList.get(rightNum)[0] != -1)) {
						setInternalpMatrix(m_nodes[internalNodeNum.get(i)]);
						System.out.println("remove element:"
								+ internalNodeNum.get(i)
								+ "and its left child is:"
								+ m_nodes[leftNum].getNr()
								+ "and its right child is:"
								+ m_nodes[rightNum].getNr());
						// it is important to update the state of this node
						// under
						// consideration, otherwise, you will get error
						this.setInternalCharState(m_nodes[internalNodeNum
								.get(i)]);
						internalNodeNum.remove(i);
					}
				}
				if (m_nodes[internalNodeNum.get(i)].getLength() == 0) {
					// get the node number of the left child node and right
					// child
					// node
					int leftNum = m_nodes[internalNodeNum.get(i)].m_left
							.getNr();
					int rightNum = m_nodes[internalNodeNum.get(i)].m_right
							.getNr();
					// if both left and right child node has been updated, then
					// update this internal node
					if ((charStateList.get(leftNum)[0] != -1)
							&& (charStateList.get(rightNum)[0] != -1)) {
						System.out.println("remove element:"
								+ internalNodeNum.get(i)
								+ "and its left child is:" + leftNum
								+ "and its right child is:" + rightNum
								+ "and its extra child is:"
								+ root.m_left.getNr());
						// it is important to update the state of this node
						// under
						// consideration, otherwise, you will get error
						this.setSudoRootCharState(m_nodes[internalNodeNum
								.get(i)]);
						internalNodeNum.remove(i);
					}
				}
			}
		}

		// trace back and set CharState for internal nodes
		traceBackCharStates(siteNum, m_nodes[root.m_right.getNr()]);

		// initialize CharStates for all nodes
		for (int i = 0; i < m_nodes.length; i++) {
			charStateList.set(i, new int[] { -1, -1, -1, -1 });
		}

	}// end of Pupko

	/* this function is used to set Char state for leaf nodes */
	public void setLeafCharState(int LeafNum, int[] LeafCharState,
			int[] m_states) {
		int LeafChar = m_states[LeafNum];
		for (int i = 0; i < LeafCharState.length; i++) {
			LeafCharState[i] = LeafChar;
		}
		charStateList.set(LeafNum, LeafCharState);
	}

	/* this func is used to set pMatrix for leaf nodes */
	public void setLeafpMatrix(int LeafChar, double[][] pMatrix,
			double branchLength) {
		for (int i = 0; i < 4; i++) {
			pMatrix[i][LeafChar] = this.transitProbJC69(i, LeafChar,
					branchLength);
		}
	}

	/* set pMatrix for internal nodes */
	public void setInternalpMatrix(Node internalNode) {
		double P_ij = 0.0;
		int leftNum = internalNode.m_left.getNr();
		int rightNum = internalNode.m_right.getNr();
		for (int i = 0; i < 4; i++) {
			for (int j = 0; j < 4; j++) {
				P_ij = this.transitProbJC69(i, j, internalNode.getLength());
				P_ij = P_ij * getpMatrixRowSum(leftNum, j)
						* getpMatrixRowSum(rightNum, j);
				setpMatrix(internalNode.getNr(), i, j, P_ij);
			}
		}
		System.out.println("The internal node is:"
				+ internalNode.getNr()
				+ "\n and the branch length is:"
				+ internalNode.getLength()
				+ " and the pMatrix is: \n"
				+ Arrays.deepToString(pMatrixList.get(internalNode.getNr()))
						.replace("], ", "\n").replaceAll(",|\\[|\\]", ""));
	}

	/* set Char State for internal nodes */
	public void setInternalCharState(Node internalNode) {
		int[] newCharState = new int[] { -1, -1, -1, -1 };
		/*
		 * this array is used to store temp probability vector of the child
		 * CharState j (A,T,C,G) given the parent CharState i
		 */
		double[] pChildStateGivenParent = new double[4];
		// loop through all ancestral states
		for (int i = 0; i < 4; i++) {
			// calculate pChildStateGivenParent
			for (int j = 0; j < 4; j++) {
				pChildStateGivenParent[j] = pMatrixList.get(internalNode
						.getNr())[i][j]
						/ getpMatrixRowSum(internalNode.getNr(), i);
			}
			// sample the newCharState for this internal node based on
			// pChildStateGivenParent
			newCharState[i] = NucleotideSampling(pChildStateGivenParent);
			System.out.println("The p vector is: "
					+ Arrays.toString(pChildStateGivenParent)
					+ " and the nucleo sampled is: " + newCharState[i]);
		}
		// set CharState for this internal node finally
		charStateList.set(internalNode.getNr(), newCharState);
	}

	/*
	 * This function might be temporary: it calculate the probability that
	 * nucleotide i at the beginning of the branch and nucleotide j at the end
	 * of the branch based on JC69 model
	 */
	public double transitProbJC69(int i, int j, double branchLength) {
		double transitProb = 0.0;
		if (i == j) {
			transitProb = 0.25 * (1 + 3 * Math.exp((-4.0 / 3.0) * branchLength));
		} else {
			transitProb = 0.25 * (1 - Math.exp((-4.0 / 3.0) * branchLength));
		}
		return transitProb;
	}

	// get the row sum of the pMatrix given the node number and row number
	public double getpMatrixRowSum(int nodeNum, int RowNum) {
		double RowSum = 0.0;
		for (int j = 0; j < 4; j++) {
			RowSum = RowSum + pMatrixList.get(nodeNum)[RowNum][j];
		}
		return RowSum;
	}

	// set the single entry in the pMatrix for a specific node
	public void setpMatrix(int nodeNum, int i, int j, double pSingleEntry) {
		pMatrixList.get(nodeNum)[i][j] = pSingleEntry;
	}

	/* set root char state */
	public void setSudoRootCharState(Node sudoRootNode) {
		// int to store the sudoRootState
		int sudoRoot = -1;

		// get the node number of the left child node and right
		// child
		// node
		int leftNum = sudoRootNode.m_left.getNr();
		int rightNum = sudoRootNode.m_right.getNr();
		int outgroupNum = root.m_left.getNr();

		double[] pChildStateGivenParent = new double[4];
		/*
		 * vector to store nucleotide frequency Here, since we are using JC69,
		 * frequencies will cancel out but if we use other model and calculate
		 * frequencies based on our data, we may set this value at the beginning
		 */
		double[] NucleoFreq = new double[] { 0.25, 0.25, 0.25, 0.25 };
		double sudoRootProbSum = 0.0;
		for (int j = 0; j < 4; j++) {
			// System.out.println("j is:"+j+" and prob is: " + NucleoFreq[j] *
			// m_nodes[leftNum].getpMatrixRowSum(j) *
			// m_nodes[rightNum].getpMatrixRowSum(j) *
			// m_nodes[outgroupNum].getpMatrixRowSum(j));
			sudoRootProbSum += NucleoFreq[j] * getpMatrixRowSum(leftNum, j)
					* getpMatrixRowSum(rightNum, j)
					* getpMatrixRowSum(outgroupNum, j);
		}
		for (int j = 0; j < 4; j++) {
			pChildStateGivenParent[j] = (NucleoFreq[j]
					* getpMatrixRowSum(leftNum, j)
					* getpMatrixRowSum(rightNum, j) * getpMatrixRowSum(
					outgroupNum, j))
					/ sudoRootProbSum;
		}
		sudoRoot = this.NucleotideSampling(pChildStateGivenParent);
		charStateList.set(sudoRootNode.getNr(), new int[] { sudoRoot, sudoRoot,
				sudoRoot, sudoRoot });
		System.out.println("The p vector is: "
				+ Arrays.toString(pChildStateGivenParent)
				+ " and the root nucleo sampled is: " + sudoRoot);
	}

	/*
	 * trace back from the root to leafs and set the Char States for all nodes
	 * at the given site
	 */
	public void traceBackCharStates(int siteNum, Node internalNode) {
		int nodeNum = internalNode.getNr();
		// deal with sudoRootNode
		if (!internalNode.isLeaf()) {
			int parentNum = internalNode.m_Parent.getNr();
			int leftNum = internalNode.m_left.getNr();
			int rightNum = internalNode.m_right.getNr();

			if (internalNode.m_Parent.isRoot()) {
				sequences[nodeNum][siteNum] = charStateList.get(nodeNum)[0];
				System.out.println("This sudoRoot Num is: "
						+ internalNode.getNr() + " and state is: "
						+ sequences[nodeNum][siteNum]);
				traceBackCharStates(siteNum, m_nodes[leftNum]);
				traceBackCharStates(siteNum, m_nodes[rightNum]);
			} else {
				// deal with other internal nodes
				sequences[nodeNum][siteNum] = charStateList.get(nodeNum)[sequences[parentNum][siteNum]];
				System.out.println("This internal Node is: "
						+ internalNode.getNr() + " and its state is: "
						+ sequences[nodeNum][siteNum]);
				traceBackCharStates(siteNum, m_nodes[leftNum]);
				traceBackCharStates(siteNum, m_nodes[rightNum]);
			}
		} else {
			System.out.println("Reach leaf: " + nodeNum);
		}
	}

	/*
	 * sample nucleotide based on discrete probability distribution this
	 * function may form its own class and move to somewhere else This function
	 * implements inverse sampling in the discrete case. For example:
	 */
	// X p(X) F(X)
	// 0 0.1 0.1
	// 1 0.3 0.4=0.1 + 0.3
	// 2 0.2 0.6=0.4+0.2
	// 3 0.4 1=0.6+0.4
	/*
	 * The first outcome for which the random deviate is smaller than (or is
	 * equal to) the associated cumulative probability corresponds to the
	 * sampled outcome
	 */

	public int NucleotideSampling(double[] pNucleotide) {
		// create an integer to store final result
		int sampledNucleo = -1;
		// from p(X) we calculate the CDF F(X)
		double[] fNucleotide = new double[4];
		fNucleotide[0] = pNucleotide[0];
		for (int i = 1; i < pNucleotide.length; i++) {
			fNucleotide[i] = fNucleotide[i - 1] + pNucleotide[i];
		}
		// draw a random number between 0 and 1
		double randNum = generator.nextDouble();
		// get the first nucleotide whose F(X) is larger or equal to randNum
		for (int i = 0; i < fNucleotide.length; i++) {
			if (fNucleotide[i] > randNum) {
				sampledNucleo = i;
				break;
			}
		}
		return sampledNucleo;
	}
}